package ink.mint.matrix;

import static java.lang.Math.floorMod;

public class MatrixExtensions extends Matrix {

    /**
     * 步长
     */
    private int stride;

    /**
     * 无参构造函数
     */
    public MatrixExtensions() {
        super();
    }

    /**
     * 构造函数
     *
     * @param row 向量的行数
     */
    public MatrixExtensions(int row) {
        super(row);
    }

    /**
     * 构造函数
     *
     * @param row    向量的行数
     * @param column 向量的列数
     */
    public MatrixExtensions(int row, int column) {
        super(row, column);
    }

    /**
     * 私有函数从大矩阵第row行col列中取出n行m列矩阵
     *
     * @param A      输入矩阵
     * @param row    第row行
     * @param column 第column列
     * @param n      取出矩阵为n行
     * @param m      取出矩阵为m列
     * @param stride 步长
     * @return 取出的矩阵
     */
    private double[][] getMinMatrix(double[][] A, int row, int column, int n, int m, int stride) {
        int ar = A.length;
        int ac = A[0].length;
        double[][] result = new double[n][m];
        int astrider = 0;
        int astridec = 0;
        for (int i = 0; i < n; i++) {
            astrider = row * stride + i;
            for (int j = 0; j < m; j++) {
                astridec = column * stride + j;
                result[i][j] = A[astrider][astridec];
            }
        }
        return result;
    }

    /**
     * 公有函数给矩阵加0边
     *
     * @param A          输入矩阵
     * @param paddingRow 行加padding
     * @param paddingCol 列加padding
     * @return 返回修正padding的矩阵
     */
    public double[][] addPaddingMatrix(double[][] A, int paddingRow, int paddingCol) {
        int ar = A.length;
        int ac = A[0].length;
        int pr = ar + paddingRow;
        int pc = ac + paddingCol;
        double[][] result = new double[pr][pc];

        for (int i = 0; i < ar; i++) {
            for (int j = 0; j < ac; j++) {
                result[i+paddingRow/2][j+paddingCol/2] = A[i][j];
            }
        }
        return result;
    }

    /**
     * 私有函数卷积操作
     *
     * @param inputA 输入矩阵A
     * @param kernel 卷积核
     * @return 卷积结果
     */
    private double convolution(double[][] inputA, double[][] kernel) {
        int ar = inputA.length;
        int ac = inputA[0].length;
        double result = 0;
        for (int i = 0; i < ar; i++) {
            for (int j = 0; j < ac; j++) {
                result += inputA[i][j] * kernel[i][j];
            }
        }
        return result;
    }

    /**
     * 私有函数计算滑块
     *
     * @param A   输入矩阵A
     * @param B   输入矩阵B
     * @param row 增加的行
     * @param col 增加的列
     * @return 加padding后的矩阵
     */
    private double[][] CalculateSlider(double[][] A, double[][] B, int row, int col) {
        int ar = A.length;
        int ac = A[0].length;
        int br = B.length;
        int bc = B[0].length;
        double[][] result = new double[ar + row * 2][ac + col * 2];
        return result = addPaddingMatrix(A, row, col);
    }

    /**
     * 私有函数取出矩阵最大值
     *
     * @param A 输入矩阵
     * @return 最大值
     */
    private double[] matrixMaxNumber(double[][] A) {
        int ar = A.length;
        int ac = A[0].length;
        double max[] = new double[3];
        double maxNumber =  Double.NEGATIVE_INFINITY;
        for (int i = 0; i < ar; i++) {
            for (int j = 0; j < ac; j++) {
                if (A[i][j] > maxNumber) {
                    max[0] = A[i][j];
                    max[1] = i;
                    max[2] = j;
                }
            }
        }
        return max;
    }

    /**
     * 私有函数取出矩阵平均值
     *
     * @param A 输入矩阵
     * @return 平均值
     */
    private double matrixAverageNumber(double[][] A) {
        int ar = A.length;
        int ac = A[0].length;
        double average = 0;
        for (int i = 0; i < ar; i++) {
            for (int j = 0; j < ac; j++) {
                average += A[i][j];
            }
        }
        average /= (ar * ac);
        return average;
    }

    /**
     * 公有函数计算特征图
     *
     * @param input  输入矩阵
     * @param kernel 卷积核
     * @param stride 步长
     * @return 计算出的特征图
     */
    public double[][] featureMap(double[][] input, double[][] kernel, int stride) {
        this.stride = stride;
        int ar = input.length;
        int ac = input[0].length;
        int br = kernel.length;
        int bc = kernel[0].length;
        int aPaddingRow = (floorMod(ar - br, stride) == 0 ? floorMod(ar - br, stride) : stride - (floorMod(ar - br, stride)));
        int aPaddingCol = (floorMod(ac - bc, stride) == 0 ? floorMod(ac - bc, stride) : stride - (floorMod(ac - bc, stride)));
        double[][] A = CalculateSlider(input, kernel, aPaddingRow, aPaddingCol);
        int Ar = A.length;
        int Ac = A[0].length;
        int maxStrideRowCount = (floorMod(Ar - br, stride) == 0 ? Math.floorDiv(Ar - br, stride) : (Math.floorDiv(Ar - br, stride) + 1)) + 1;
        int maxStrideColConut = (floorMod(Ac - bc, stride) == 0 ? Math.floorDiv(Ac - bc, stride) : (Math.floorDiv(Ac - bc, stride) + 1)) + 1;
        double[][] featureMap = new double[maxStrideRowCount][maxStrideColConut];
        for (int i = 0; i < maxStrideRowCount; i++) {
            for (int j = 0; j < maxStrideColConut; j++) {
                featureMap[i][j] = convolution(getMinMatrix(A, i, j, br, bc, stride), kernel);
            }
        }
        return featureMap;
    }

    public double[][] featureMap(double[][] input, double[][] kernel, double bias, int stride, int padding) {
        this.stride = stride;
        int br = kernel.length;
        int bc = kernel[0].length;
        double[][] A = CalculateSlider(input, kernel, padding, padding);
        int Ar = A.length;
        int Ac = A[0].length;
        int maxStrideRowCount = (Ar - br) / stride + 1;
        int maxStrideColCount = (Ac - bc) / stride + 1;
        double[][] featureMap = new double[maxStrideRowCount][maxStrideColCount];
        for (int i = 0; i < maxStrideRowCount; i++) {
            for (int j = 0; j < maxStrideColCount; j++) {
                featureMap[i][j] = convolution(getMinMatrix(A, i, j, br, bc, stride), kernel) + bias;

            }
        }
//        System.out.println(" feature  map:  "+featureMap.length+" | "+featureMap[0].length);
        return featureMap;
    }

    public double[][] featureMap(double[][] input, double[][] kernel, int stride, int padding) {
        this.stride = stride;
        int br = kernel.length;
        int bc = kernel[0].length;
        double[][] A = CalculateSlider(input, kernel, padding, padding);
        int Ar = A.length;
        int Ac = A[0].length;
        int maxStrideRowCount = (Ar - br) / stride + 1;
        int maxStrideColConut = (Ac - bc) / stride + 1;
        double[][] featureMap = new double[maxStrideRowCount][maxStrideColConut];
        for (int i = 0; i < maxStrideRowCount; i++) {
            for (int j = 0; j < maxStrideColConut; j++) {
                featureMap[i][j] = convolution(getMinMatrix(A, i, j, br, bc, stride), kernel);
//                System.out.print(" feature  map:  "+featureMap[i][j]+" | ");

            }
//            System.out.println();
        }
        return featureMap;
    }


    /**
     * 公有函数最大化池化
     *
     * @param input 输入矩阵
     * @param size  池化大小（正方形）
     * @return 最大化池化矩阵
     */
    public double[][] maxPooling(double[][] input, int size) {
        int inputr = input.length;
        int inputc = input[0].length;
        double[][] A = input;
        if (floorMod(inputr, size) > 0) {
            int paddingRow = size - floorMod(inputr, size);
            A = addPaddingMatrix(A, paddingRow, 0);
        }
        if (floorMod(inputc, size) > 0) {
            int paddingCol = size - floorMod(inputc, size);
            A = addPaddingMatrix(A, 0, paddingCol);
        }
        int ar = A.length;
        int ac = A[0].length;
        double[][] result = new double[Math.floorDiv(ar, size)][Math.floorDiv(ac, size)];
        int rr = result.length;
        int rc = result[0].length;
        for (int i = 0; i < rr; i++) {
            for (int j = 0; j < rc; j++) {
                double[] matrixMaxNumbers = matrixMaxNumber(getMinMatrix(A, i, j, size, size, stride));
                result[i][j] = matrixMaxNumbers[0];
            }
        }
        return result;
    }

    int[][] maxIndex;

    public double[][] maxPooling(double[][] input, int size, int stride, int padding) {
        int inputr = input.length;
        int inputc = input[0].length;
        double[][] A = input;
        if (floorMod(inputr, size) > 0) {
            A = addPaddingMatrix(A, padding, 0);
        }
        if (floorMod(inputc, size) > 0) {
            A = addPaddingMatrix(A, 0, padding);
        }
        int ar = A.length;
        int ac = A[0].length;
        double[][] result = new double[Math.floorDiv(ar, size)][Math.floorDiv(ac, size)];
        int rr = result.length;
        int rc = result[0].length;
        maxIndex = new int[rr * rc][2];
        int index = 0;
        for (int i = 0; i < rr; i++) {
            for (int j = 0; j < rc; j++) {
                double[] matrixMaxNumbers = matrixMaxNumber(getMinMatrix(A, i, j, size, size, stride));
                result[i][j] = matrixMaxNumbers[0];
//                System.out.print(" i: " + i + " | j: " + j);
//                System.out.println();
//                if (result[i][j] == 0) {
//                    maxIndex[index][0] = 0;
//                    maxIndex[index][1] = 0;
//                } else {
                    maxIndex[index][0] = (int)matrixMaxNumbers[1] + i * stride;
                    maxIndex[index][1] = (int)matrixMaxNumbers[2] + j * stride;
//                }
//                System.out.print("maxpooling index: " + index + " i:" + maxIndex[index][0] + " | j:" + maxIndex[index][1]+"  ||  "+result[i][j]);
                index++;

            }
//            System.out.println();
        }
//        System.out.println("maxpooling   index: "+maxIndex.length);
        return result;
    }


    /**
     * 公有方法均分池化
     *
     * @param input 输入矩阵
     * @param size  均分池化块大小
     * @return 均分池化矩阵
     */
    public double[][] averagePooling(double[][] input, int size) {
        int inputr = input.length;
        int inputc = input[0].length;
        double[][] A = input;
        if (floorMod(inputr, size) > 0) {
            int paddingRow = size - floorMod(inputr, size);
            A = addPaddingMatrix(A, paddingRow, 0);
        }
        if (floorMod(inputc, size) > 0) {
            int paddingCol = size - floorMod(inputc, size);
            A = addPaddingMatrix(A, 0, paddingCol);
        }
        int ar = A.length;
        int ac = A[0].length;
        double[][] result = new double[Math.floorDiv(ar, size)][Math.floorDiv(ac, size)];
        int rr = result.length;
        int rc = result[0].length;
        for (int i = 0; i < rr; i++) {
            for (int j = 0; j < rc; j++) {
                result[i][j] = matrixAverageNumber(getMinMatrix(A, i, j, size, size, size));
            }
        }
        return result;
    }

    public double[][] averagePooling(double[][] input, int size, int stride, int padding) {
        int inputr = input.length;
        int inputc = input[0].length;
        double[][] A = input;
        if (floorMod(inputr, size) > 0) {
            A = addPaddingMatrix(A, padding, 0);
        }
        if (floorMod(inputc, size) > 0) {
            A = addPaddingMatrix(A, 0, padding);
        }
        int ar = A.length;
        int ac = A[0].length;
        double[][] result = new double[Math.floorDiv(ar, size)][Math.floorDiv(ac, size)];
        int rr = result.length;
        int rc = result[0].length;
        for (int i = 0; i < rr; i++) {
            for (int j = 0; j < rc; j++) {
                result[i][j] = matrixAverageNumber(getMinMatrix(A, i, j, size, size, stride));
            }
        }
        return result;
    }


    /**
     * 创建随机核函数
     *
     * @param kernelSize 核大小
     * @return 一个kernelSize X kernelSize 的随机核函数
     */
    public double[][] createRandomKernel(int kernelSize) {
        return random(kernelSize, kernelSize);
    }


    public double[] matrixToVector(double[][] A) {
        int ar = A.length;
        int ac = A[0].length;
        int index = 0;
        double[] result = new double[ar * ac];
        for (int i = 0; i < ar; i++) {
            for (int j = 0; j < ac; j++) {
                result[index++] = A[i][j];
            }
        }
        return result;
    }

    public double[][] vectorDismember(double[] A, int dismember, int minMatrixSize) {
        int ar = A.length;
        double[][] minMatrix = new double[dismember][minMatrixSize * minMatrixSize];
        int index = 0;
        for (int i = 0; i < dismember; i++) {
            for (int j = 0; j < minMatrixSize * minMatrixSize; j++) {
                minMatrix[i][j] = A[index++];
            }
        }
        return minMatrix;
    }

    public double[][] rot180Matrix(double[][] A) {
        int ar = A.length;
        int ac = A[0].length;
        double[][] result = new double[ar][ac];
        for (int i = 0; i < ar; i++) {
            for (int j = 0; j < ac; j++) {
                result[i][j] = A[ar - 1 - i][ac - 1 - j];
            }
        }
        return result;
    }

    public int[][] getMaxIndex() {
        return maxIndex;
    }

    public void setMaxIndex(int[][] maxIndex) {
        this.maxIndex = maxIndex;
    }

    // MatrixExtensions.java 完善反卷积方法
//    public double[][] deconvolution(double[][] delta,
//                                    double[][] rotKernel,
//                                    int stride,
//                                    int padding,
//                                    int targetSize) {
//        // 计算全卷积尺寸
//        int outputSize = (delta.length - 1) * stride + rotKernel.length - 2 * padding;
//        if (outputSize != targetSize) {
//            throw new IllegalArgumentException(
//                    "目标尺寸计算错误: 预期 " + targetSize +
//                            " 实际 " + outputSize +
//                            "\n输入尺寸: " + delta.length +
//                            "\n核尺寸: " + rotKernel.length +
//                            "\n步长: " + stride +
//                            "\n填充: " + padding
//            );
//        }
//
//        // 创建输出矩阵
//        double[][] result = new double[outputSize][outputSize];
//
//        // 执行全卷积操作
//        for (int i = 0; i < delta.length; i++) {
//            for (int j = 0; j < delta[0].length; j++) {
//                // 计算起始位置
//                int xStart = i * stride;
//                int yStart = j * stride;
//
//                // 叠加核的贡献
//                for (int k = 0; k < rotKernel.length; k++) {
//                    for (int l = 0; l < rotKernel[0].length; l++) {
//                        int x = xStart + k - padding;
//                        int y = yStart + l - padding;
//
//                        if (x >= 0 && x < outputSize && y >= 0 && y < outputSize) {
//                            result[x][y] += delta[i][j] * rotKernel[k][l];
//                        }
//                    }
//                }
//            }
//        }
//        return result;
//    }

    public double[][] deconvolution(double[][] delta,
                                    double[][] rotKernel,
                                    int stride,
                                    int padding,
                                    int targetSize) {
        // 参数校验
        if (delta == null || rotKernel == null) {
            throw new IllegalArgumentException("Input matrices cannot be null");
        }
        if (stride <= 0 || padding < 0 || targetSize <= 0) {
            throw new IllegalArgumentException("Invalid parameters: stride=" + stride
                    + ", padding=" + padding + ", targetSize=" + targetSize);
        }

        // 计算预期输出尺寸
        int kernelSize = rotKernel.length;
        int calculatedSize = (delta.length - 1) * stride + kernelSize - 2 * padding;

        // 尺寸验证
        if (calculatedSize != targetSize) {
            throw new IllegalArgumentException(
                    "Target size mismatch!\n" +
                            "Calculated size: " + calculatedSize + "\n" +
                            "Requested size: " + targetSize + "\n" +
                            "Formula: (delta_length-1)*stride + kernel_size - 2*padding\n" +
                            "Delta dimensions: " + delta.length + "x" + delta[0].length + "\n" +
                            "Kernel dimensions: " + kernelSize + "x" + kernelSize + "\n" +
                            "Stride: " + stride + ", Padding: " + padding
            );
        }

        // 初始化输出矩阵
        double[][] output = new double[targetSize][targetSize];

        // 遍历delta矩阵的每个元素
        for (int i = 0; i < delta.length; i++) {
            for (int j = 0; j < delta[0].length; j++) {

                // 计算对应的起始位置（考虑padding）
                int xStart = i * stride - padding;
                int yStart = j * stride - padding;

                // 应用旋转后的卷积核
                for (int k = 0; k < kernelSize; k++) {
                    for (int l = 0; l < kernelSize; l++) {
                        // 计算实际坐标
                        int x = xStart + k;
                        int y = yStart + l;

                        // 边界检查
                        if (x >= 0 && x < targetSize && y >= 0 && y < targetSize) {
                            // 累加卷积结果
                            output[x][y] += delta[i][j] * rotKernel[k][l];
                        }
                    }
                }
            }
        }
        return output;
    }

    public double[][] validConvolution(double[][] input, double[][] kernel, int stride, int padding) {
        // 参数校验
        if (input == null || kernel == null) {
            throw new IllegalArgumentException("Input matrices cannot be null");
        }
        if (stride <= 0) {
            throw new IllegalArgumentException("Stride must be positive");
        }

        int inputSize = input.length;
        int kernelSize = kernel.length;

        // 验证输入尺寸
        if (inputSize == 0 || kernelSize == 0 || kernelSize > inputSize) {
            throw new IllegalArgumentException(
                    "Invalid dimensions: input=" + inputSize + "x" + inputSize
                            + ", kernel=" + kernelSize + "x" + kernelSize
            );
        }

        // 计算输出尺寸
        int outputSize = (inputSize - kernelSize) / stride + 1;
        double[][] output = new double[outputSize][outputSize];

        // 执行有效卷积
        for (int i = 0; i <= inputSize - kernelSize; i += stride) {
            for (int j = 0; j <= inputSize - kernelSize; j += stride) {

                // 计算当前窗口的卷积值
                double sum = 0.0;
                for (int k = 0; k < kernelSize; k++) {
                    for (int l = 0; l < kernelSize; l++) {
                        sum += input[i + k][j + l] * kernel[k][l];
                    }
                }

                // 存储结果
                output[i/stride][j/stride] = sum;
            }
        }
        return output;
    }

}
